
# Copyright (C), 2020,WHU.
#   File name:functions.py      /
#   Created by:
#   Version:1.0.0        
#   Date: 
#   Description:  实现lenet网络中需要用到的激活函数，损失函数，卷积，池化，以及对应的反向传播的函数
#   
# Dependencies：

# ChangeLog:        // 修改历史记录列表，每条修改记录应包括修改日期、修改
#                   // 者及修改内容简述  

import numpy as np


#relu函数实现
def relu(Z):
     return Z*(Z > 0)

#relu导函数实现
def d_relu(A):
    return np.ones_like(A)*(A > 0)
#sigmoid函数实现
def sigmoid(Z):
    return 1/(1+np.exp(-Z))

#d_sigmoid导函数实现
def d_sigmoid(A):
    return A*(1-A)

#softmax函数实现
def softmax(Z):
    return np.exp(Z)/np.sum(np.exp(Z), axis=0, keepdims=True)
#定义交叉熵损失函数
def ELoss(y, y_hat):
    return np.sum(-y * np.log(y_hat))/y.shape[1]

#提高卷积效率，转化成行列相乘的卷积处理方式
# def im2col(img, ksize, stride=1):
#     return 
#根据步长分割输入的对象
def split_by_strides(X, k_h, k_w, s):
    #例子，32*32的图像输入，s=1，k_h,K_w=5，用5*5卷积核卷积图像得到28*28的输出
    #return A,the size of A (batchsize,28,28,5,5,1)
    n, w, h, c = X.shape
    o_h = (h - k_h) // s + 1
    o_w = (w - k_w) // s + 1
    shape = (n, o_h, o_w, k_h, k_w, c)
    strides = (X.strides[0], X.strides[1] * s, X.strides[2] * s, *X.strides[1:])#strides是numpy数组对象的一个属性，官方手册给出的解释是跨越数组各个维度所需要经过的字节数（bytes）
    A = np.lib.stride_tricks.as_strided(X, shape=shape, strides=strides)
    return A
#定义卷积函数
def conv(X, W, b, s=1, padding=None):
    #调用split_by_strides
    k_w, k_h, c_in, c_out = W.shape
    if padding:
        X = np.pad(X, ((0, 0), (padding, padding), (padding, padding), (0, 0)), 'constant')
    X_split = split_by_strides(X, k_w, k_h, s)
    feature_map = np.tensordot(X_split, W, axes=[(3, 4, 5), (0, 1, 2)]) + b
    return feature_map

#定义池化函数
def pool(A, kernel_size=2, stride=2, padding=None, pool_mode='max'):
#均值池化和最大池化两种方式
    if padding:
        A = np.pad(A, padding, mode='constant')
    A_w = split_by_strides(A, kernel_size, kernel_size, stride)

    if pool_mode == 'max':
        return A_w.max(axis=(3, 4))
    elif pool_mode == 'avg':
        return A_w.mean(axis=(3, 4))
#定义转置卷积
def d_conv(dZ, A, W, s=1):
    m, k_w, k_h, k_num = dZ.shape
    A_split = split_by_strides(A, k_w, k_h, s)
    dW = np.tensordot(A_split, dZ, axes=[(0, 3, 4), (0, 1, 2)])/m
    db = np.sum(dZ, axis=(0, 1, 2)).reshape(1, -1)/m
    W = W[::-1, ::-1]
    W = W.transpose((0, 1, 3, 2))
    dA = conv(dZ, W, 0, padding=4)
    return dW, db, dA